# -*- coding: utf-8 -*-
import cv2;
import math;
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

def DarkChannel(im,sz):#计算图像的暗通道
    b,g,r = cv2.split(im)
    dc = cv2.min(cv2.min(r,g),b)#求两次最小值操作
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(sz,sz))#利用方框滤波
    dark = cv2.erode(dc,kernel)
    return dark

def AtmLight(im,dark):#计算大气光A
    [h,w] = im.shape[:2]

    imsz = h*w
    numpx = int(max(math.floor(imsz/1000),1))#前0.1%个像素点的个数
    darkvec = dark.reshape(imsz,1)#将暗通道序列化
    imvec = im.reshape(imsz,3);#将图像序列化
    indices = np.argsort(darkvec,0)#将暗通道按灰度值排序
    indices = indices[imsz-numpx::]#取灰度值最大的的前0.1%
    b,g,r=cv2.split(im)
    gray_im=r*0.299 + g*0.587 + b*0.114#计算彩色图像对应的灰度值
    gray_im=gray_im.reshape(imsz,1)
    loc=np.where(gray_im==max(gray_im[indices]))#找出暗通道中的前0.1%对应原始彩色图像中灰度值最大的点
    x=loc[0][0]
    A=np.array(imvec[x])
    A=A.reshape(1,3)
    return A#返回该点灰度值，作为大气光


def TransmissionEstimate(im,A,sz):#进行传输率的估计
    omega = 1;#控制天空的影响程度
    im3 = np.empty(im.shape,im.dtype)

    for ind in range(0,3):
        im3[:,:,ind] = im[:,:,ind]/A[0,ind]#每个通道除以各自的大气光值，normalize

    transmission = 1 - omega*DarkChannel(im3,sz)
    return transmission

def Guidedfilter(im,p,r,eps):#导向滤波
    
    mean_I = cv2.boxFilter(im,cv2.CV_64F,(r,r))#对引导图像进行方框滤波，分为一个个window
    
    mean_p = cv2.boxFilter(p, cv2.CV_64F,(r,r))#对输入图像进行方框滤波
    mean_Ip = cv2.boxFilter(im*p,cv2.CV_64F,(r,r))#利用公式进行计算
    cov_Ip = mean_Ip - mean_I*mean_p

    mean_II = cv2.boxFilter(im*im,cv2.CV_64F,(r,r))
    var_I   = mean_II - mean_I*mean_I

    a = cov_Ip/(var_I + eps)
    b = mean_p - a*mean_I

    mean_a = cv2.boxFilter(a,cv2.CV_64F,(r,r))#求得a的均值矩阵
    mean_b = cv2.boxFilter(b,cv2.CV_64F,(r,r))#求得b的均值矩阵

    q = mean_a*im + mean_b#线性变换
    return q

def TransmissionRefine(im,et):#传输率细化
    gray = cv2.cvtColor(im,cv2.COLOR_BGR2GRAY)
    gray = np.float64(gray)/255
    r = 50
    eps = 0.0001
    t = Guidedfilter(gray,et,r,eps)

    return t

def Recover(im,t,A,tx = 0.1):#图像去雾，原始场景恢复
    res = np.empty(im.shape,im.dtype)
    t = cv2.max(t,tx)
    for ind in range(0,3):
        res[:,:,ind] = (im[:,:,ind]-A[0,ind])/t + A[0,ind]

    return res

if __name__ == '__main__':
    import sys
    try:
        fn = sys.argv[1]
    except:
        fn = 'images/5.jpg'

    def nothing(*argv):
        pass

    src = cv2.imread(fn)#读取图像
    I = src.astype('float64')/255#图像归一化
 
    dark = DarkChannel(I,15)#计算暗通道，patch大小为15*15
    A = AtmLight(I,dark)#计算大气光
    te = TransmissionEstimate(I,A,15)#大致估计传输率
    t = TransmissionRefine(src,te)#传输率细化
    J = Recover(I,t,A,0.1)#图像恢复
    J[J<0]=0#把恢复图像中灰度值小于0的部分置0

    cv2.imshow("dark",dark)#显示暗通道
    cv2.imshow("t",t)#传输率
    cv2.imshow('I',src)#原始图像
    cv2.imshow('J',J)#去雾图像

    cv2.waitKey(0)
    
